package org.atlantis.evnlib.internal;

import java.lang.ref.ReferenceQueue;
import java.lang.ref.WeakReference;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;

// tiny cache
public class WeakCache<K, V> {

	private ConcurrentMap<K, ValueRef<K, V>> internalMap = new ConcurrentHashMap<K, ValueRef<K, V>>();
	private ReferenceQueue<V> clearQueue = new ReferenceQueue<V>();

	public void addCache(K key, V val) {
		ValueRef<K, V> ref = new ValueRef<K, V>(key, val, clearQueue);
		internalMap.putIfAbsent(key, ref);
	}

	public V getCache(K key) {
		clearCache();
		ValueRef<K, V> ref = internalMap.get(key);
		if (ref != null) return ref.get();
		return null;
	}

	public int size() {
		clearCache();
		return internalMap.size();
	}

	public boolean exists(K key) {
		if (internalMap.containsKey(key) && internalMap.get(key).get() == null) internalMap.remove(key);
		clearCache();
		return internalMap.containsKey(key);
	}

	public void clear() {
		internalMap.clear();
	}

	@SuppressWarnings("unchecked")
	private void clearCache() {
		ValueRef<K, V> ref = null;
		while ((ref = (ValueRef<K, V>) clearQueue.poll()) != null) {
			if (ref.get() == null) {
				internalMap.remove(ref.key);
			}
		}
	}

	static class ValueRef<K, V> extends WeakReference<V> {

		final K key;

		public ValueRef(K key, V value, ReferenceQueue<V> q) {
			super(value, q);
			this.key = key;
		}
	}
}
